from configs.model_config import min_freq


class Vocab:
    bos_token = "<bos>"
    eos_token = "<eos>"
    unk_token = "<unk>"
    pad_token = "<pad>"
    bos_id = 0
    eos_id = 1
    pad_id = 2
    unk_id = 3

    def __init__(self):
        self.word2index = {
            self.bos_token: self.bos_id,
            self.eos_token: self.eos_id,
            self.pad_token: self.pad_id,
            self.unk_token: self.unk_id
        }
        # self.word2count = {}
        self.index2word = {
            self.bos_id: self.bos_token,
            self.eos_id: self.eos_token,
            self.pad_id: self.pad_token,
            self.unk_id: self.unk_token
        }
        self.n_words = 4

    def itow(self, index: int):
        return self.index2word[index]

    def wtoi(self, word: str):
        return self.word2index.get(word, self.unk_id)

    def add_word(self, word):
        if word not in self.word2index:
            self.word2index[word] = self.n_words
            # self.word2count[word] = 1
            self.index2word[self.n_words] = word
            self.n_words += 1
        else:
            pass

    def words2indices(self, words, add_special=True) -> list:
        if add_special:
            return [self.bos_id] + [self.wtoi(word) for word in words] + [self.eos_id]
        else:
            return [self.wtoi(word) for word in words]

    def indices2words(self, indices):
        return [self.itow(index) for index in indices]

    def __len__(self):
        return len(self.index2word)


def build_vocab_from_file(file_path) -> Vocab:
    with open(file_path) as file:
        vocab = Vocab()
        lines = file.readlines()
        for line in lines:
            word, count = line.split()
            word, count = word, int(count)
            if count > min_freq:
                vocab.add_word(word)
        return vocab
